Skip to content

feat(executorch): expose composable Edge export API - #4440

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:executorch-composable-export
Open

feat(executorch): expose composable Edge export API#4440
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:executorch-composable-export

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Why this PR is needed

torch_tensorrt.save(..., output_format="executorch") is the easiest way to create an ExecuTorch .pte file. It completes the whole export pipeline immediately:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn(1, 3, 224, 224, device="cuda")]
trt_module = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)

torch_tensorrt.save(
    trt_module,
    "model.pte",
    output_format="executorch",
)

That remains the recommended API for common one-step exports.

Some workflows need to stop before final ExecuTorch memory planning and serialization. For example, users may want to:

  • inspect the delegated Edge graph;
  • run custom Edge transforms;
  • send operations that TensorRT does not support to another backend;
  • preserve independent methods such as prefill and decode;
  • add constant methods or generate an ETRecord;
  • choose the final ExecuTorch backend configuration.

This PR adds torch_tensorrt.executorch.export(). It returns ExecuTorch's standard EdgeProgramManager, which is the supported boundary for inspecting and customizing an Edge program before calling to_executorch().

Which API should I use?

Use torch_tensorrt.save() when you want a .pte in one step.

Use torch_tensorrt.executorch.export() when you need control before the final .pte is created:

import torch_tensorrt.executorch

edge = torch_tensorrt.executorch.export(trt_module)

# Inspect or transform the Edge program here.
print(edge.exported_program().graph)

program = edge.to_executorch()
with open("model.pte", "wb") as output:
    program.write_to_file(output)

# Advanced callers must also persist external tensor data when present.
program.write_tensor_data_to_file(".")

save() now delegates to the same implementation, so the simple and advanced paths do not maintain separate lowering logic.

Example: TensorRT with a fallback backend

TensorRT always receives the first chance to claim its prebuilt engine nodes. Additional partitioners run afterward and can claim operations that TensorRT does not support.

Assume cuda_partitioner is an ExecuTorch partitioner configured for your target:

edge = torch_tensorrt.executorch.export(
    trt_module,
    partitioners=[cuda_partitioner],
)

program = edge.to_executorch()

This can produce one ExecuTorch program containing both TensorRT and CUDA delegates.

Example: multiple methods

Assume prefill_program and decode_program are independent, engine-bearing ExportedProgram objects, and the two spec lists are configured for their methods:

edge = torch_tensorrt.executorch.export(
    {
        "prefill": prefill_program,
        "decode": decode_program,
    },
    compile_specs={
        "prefill": prefill_specs,
        "decode": decode_specs,
    },
)

program = edge.to_executorch()

The mapping preserves method names and allows per-method partitioner and compile-spec pipelines. It does not imply that mutable state is shared between methods.

Example: transforms, metadata, and ETRecord

Assume the transform passes and Edge configuration below are provided by the caller:

edge = torch_tensorrt.executorch.export(
    trt_module,
    transform_passes=my_edge_passes,
    compile_config=edge_compile_config,
    constant_methods={"get_vocab_size": 256},
    generate_etrecord=True,
)

program = edge.to_executorch()
program.get_etrecord().save("model_etrecord.bin")

Accepted inputs

The new API accepts:

  • a TensorRT-compiled torch.fx.GraphModule;
  • an engine-bearing torch.export.ExportedProgram;
  • a mapping from method names to independent ExportedProgram objects.

A plain torch.nn.Module must be compiled with Torch-TensorRT first.

Safety and memory behavior

Engine rewriting changes the exported graph. The new API stages an independent copy of graph structure, signatures, state containers, constants, and metadata before rewriting anything. Rewrite and lowering failures therefore do not consume or structurally modify the caller's original program.

Tensor and TensorRT engine payload storage is treated as immutable and shared with the staged structure. This avoids deep-copying potentially multi-gigabyte engines. Custom transform passes must not mutate shared payload objects.

Metadata describing symbolic shapes is also shared rather than copied. It points back into the live shape environment the exported program is guarded on, and that environment holds fake tensors which cannot be copied. Sharing it is what keeps dynamic-shape models working; copying it would also detach the copy from the shape symbols the graph depends on.

When a lifted TensorRT engine is replaced, the old placeholder, graph-signature entry, and constant are removed. Multiple calls that share one engine reuse one materialized uint8 payload buffer.

Compatibility

  • Existing one-step save(..., output_format="executorch") behavior is preserved.
  • .pte serialization and external .ptd persistence are preserved.
  • TensorRT remains the first partitioner.
  • Shared dynamic dimensions declared with Input(shared_dims=...) remain shared.
  • Zero-engine methods are allowed. A later partitioner may claim them, or portable operators may remain undelegated.

Validation

Built from source with the Torch-TensorRT C++ runtime active, against TensorRT and
ExecuTorch, and exercised on both x86_64 and aarch64 CUDA GPUs.

Checked end to end with a real compiled TensorRT engine:

  • static-shape export() writes a loadable .pte;
  • dynamic-shape export() writes a loadable .pte of the same size the previous
    one-step save path produced;
  • the caller's GraphModule still holds its original engine nodes after export();
  • engine payloads are shared rather than re-serialized;
  • save(..., output_format="executorch") still writes a .pte;
  • save() still raises TypeError for a non-list partitioners or compile_specs.

Unit and integration tests in the ExecuTorch test directory: 76 passed.

Two GPU tests in test_cuda_partitioner_composition.py,
test_erfinv_routes_to_cuda_backend and
test_weighted_partition_persists_external_data, fail with
SpecViolationError: Node.meta lowered_module_0 is missing val field. They fail the
same way on main without this change, so they are pre-existing and not caused by
this work.

Note that these tests only run in the L2 tier, which requires the ci: run-l2 label
on a pull request.

Static checks on the changed files: import ordering, formatting, lint, and type
checking.

Related work

PR #4433 forwards additional Edge-lowering options through one-step save(). After this PR lands, #4433 can keep save() as a thin wrapper over torch_tensorrt.executorch.export() instead of extending the previous private lowering path.

graph_module: torch.fx.GraphModule,
payload_memo: dict[int, Any],
) -> torch.fx.GraphModule:
staged = copy.deepcopy(graph_module, payload_memo)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this doubles the GPU/Host memory usage? Is the copy absolutely necessary?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Or is there a way to not copy the weight, only the graph structure

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question to push on, and no, weights are not copied. Only graph structure is,
which is what you suggest in your next comment.

_payload_sharing_memo (just above) pre-seeds copy.deepcopy's memo with every tensor
and engine payload reachable from the program, so deepcopy returns those objects
as-is instead of duplicating them. state_dict and constants are then shallow-copied
dicts, so the keys are independent but the values are the same objects.

I measured it to be sure, on a model with 33.6 MB of weights:

weights in program: 33.6 MB
live CPU tensor bytes before staging: 33.6 MB
live CPU tensor bytes after  staging: 33.6 MB
growth: 0.0 MB

state_dict entries: 4, shared by identity: 4
  a.weight   same object: True  same storage: True
  a.bias     same object: True  same storage: True
  b.weight   same object: True  same storage: True
  b.bias     same object: True  same storage: True

So the copy is graph nodes, signatures, and metadata containers, all small. TensorRT
engine blobs are shared the same way, which matters more since those can be multiple
gigabytes.

On whether the copy is necessary at all: replace_execute_engine rewrites the graph in
place (registers a buffer, erases nodes, edits the graph signature). Without staging, a
failure part-way through would leave the caller's ExportedProgram half-rewritten. The
GraphModule path happens to be safe already because dynamo/_exporter.py deep-copies,
but the pre-exported ExportedProgram and multi-method paths are handed to us directly,
so staging is what keeps those inputs intact.

The cost this leaves is documented as a constraint: transform passes must not mutate the
shared payloads in place, since those are the caller's tensors.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that is exactly what it does. Only graph structure is copied; tensors and engine
payloads are shared by identity through the pre-seeded deepcopy memo. Measurements in
the reply above: 33.6 MB of weights in, 0.0 MB growth, all four state_dict tensors
identical objects sharing the same storage.

One case needed a carve-out, which is worth mentioning since it was a real bug. Metadata
describing symbolic shapes (SymInt, fake tensors) cannot be deep-copied at all: it
reaches back into the live ShapeEnv, which holds fake tensors, and copying raised

RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor)

on every dynamic-shape model. That metadata is now shared rather than copied, which also
keeps the staged graph attached to the shape symbols it is guarded on. Covered by
test_stage_exported_program_supports_dynamic_shapes.

graph_module.register_buffer(buffer_name, engine_tensor, persistent=True)
exported_program.state_dict[buffer_name] = engine_tensor
str_args = [
str(value) if value is not None else ""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

engine_info[ENGINE_IDX] is raw bytes (often hundreds of MB / multi-GB). str(b"...") builds a temporary Python string of roughly 4 * len(bytes) characters (b'\xab\xcd...' style), then immediately throws that slot away and uses engine_attr_node instead.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, this was a real waste and it is fixed now.

You are right that the slot is thrown away: no_op_args splices in engine_attr_node at
ENGINE_IDX, so the stringified engine was built and then never used. Measured on a
64 MB stand-in payload:

engine bytes      : 67.1 MB
str(bytes) length : 268.4 M chars
str object bytes  : 268.4 MB
blowup factor     : 4.0 x

So a 2 GB engine burned about 8 GB of host memory building a string that was immediately
discarded.

Now that index is skipped:

str_args = [
    ("" if index == ENGINE_IDX else str(value) if value is not None else "")
    for index, value in enumerate(engine_info[:SERIALIZATION_LEN])
]

Re-verified on a GPU afterward: static-shape and dynamic-shape export() still produce
loadable .pte files of the same size as before, and save(output_format="executorch")
is unchanged.

shared mutable state between them.
"""
from torch_tensorrt._features import ENABLED_FEATURES

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need a check that the platform is linux?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, we do, and it was missing. Added.

save(output_format="executorch") already rejects non-Linux platforms in _compile.py,
so the advanced entry point was skipping a check the one-step path enforces. export()
now raises the same error with the same wording:

if platform.system() != "Linux":
    raise ValueError(
        f"The executorch format is only supported on Linux, {platform.system()} "
        "is not a supported platform for this format"
    )

Placed before the runtime-feature check so the platform message comes first, and covered
by test_export_rejects_non_linux_platform.

## What this adds

`torch_tensorrt.save(..., output_format="executorch")` turns a compiled model
into an ExecuTorch `.pte` file in one step. That is the easiest path and it stays
the recommendation.

Some workflows need to stop earlier. You may want to look at the delegated Edge
graph, run your own Edge transforms, send operations TensorRT cannot handle to
another backend, keep separate methods such as `prefill` and `decode`, add
constant methods, or pick the final ExecuTorch configuration yourself.

This adds `torch_tensorrt.executorch.export()`. It returns ExecuTorch's standard
`EdgeProgramManager`, which is the supported place to inspect and customize an
Edge program before you call `to_executorch()`.

```python
import torch_tensorrt.executorch

edge = torch_tensorrt.executorch.export(trt_module)

# Inspect or transform the Edge program here.
print(edge.exported_program().graph)

program = edge.to_executorch()
with open("model.pte", "wb") as output:
    program.write_to_file(output)

# Advanced callers must also persist external tensor data when present.
program.write_tensor_data_to_file(".")
```

`save()` now calls the same implementation, so the simple and advanced paths
share one lowering path instead of two.

## Which one should I use?

Use `save()` when you just want a `.pte`. Use
`torch_tensorrt.executorch.export()` when you need control before the `.pte` is
created.

## Example: TensorRT with a fallback backend

TensorRT always gets the first chance to claim its prebuilt engine nodes. Any
partitioners you pass run afterward and can claim operations TensorRT does not
support.

```python
edge = torch_tensorrt.executorch.export(
    trt_module,
    partitioners=[cuda_partitioner],
)

program = edge.to_executorch()
```

## Example: multiple methods

```python
edge = torch_tensorrt.executorch.export(
    {
        "prefill": prefill_program,
        "decode": decode_program,
    },
    compile_specs={
        "prefill": prefill_specs,
        "decode": decode_specs,
    },
)
```

Method names are preserved, and each method can have its own partitioner and
compile-spec pipeline. This does not make mutable state shared between methods.

## Accepted inputs

- a TensorRT-compiled `torch.fx.GraphModule`
- an engine-bearing `torch.export.ExportedProgram`
- a mapping from method names to independent `ExportedProgram` objects

A plain `torch.nn.Module` must be compiled with Torch-TensorRT first.

## Safety and memory behavior

Rewriting engine calls changes the exported graph. To avoid damaging the program
you passed in, export first stages its own copy of the graph structure,
signatures, state containers, constants, and metadata. If rewriting or lowering
fails, your original program is left alone.

Tensor and TensorRT engine payloads are treated as read-only and shared with
that staged copy, so a multi-gigabyte engine is not duplicated. Custom transform
passes must not modify these shared payload objects.

Metadata that describes symbolic shapes is a special case. It points back to the
live shape environment that the exported program is guarded on, so it is shared
rather than copied. Copying it would both fail (the shape environment holds fake
tensors that cannot be duplicated) and detach the copy from the symbols the graph
depends on. This is what keeps dynamic-shape models working.

When a lifted TensorRT engine is replaced, the old placeholder, its
graph-signature entry, and its constant are removed. Several calls that share one
engine reuse a single materialized `uint8` payload buffer.

## Compatibility

- Existing one-step `save(..., output_format="executorch")` behavior is preserved,
  including its rejection of non-list `partitioners` and `compile_specs`.
- `"executorch"` is now included in the `output_format` type hint. It was already
  accepted at run time, so this only fixes the annotation.
- `.pte` serialization and external `.ptd` persistence are preserved.
- TensorRT remains the first partitioner.
- Shared dynamic dimensions declared with `Input(shared_dims=...)` remain shared.
- Zero-engine methods are allowed. A later partitioner may claim them, or portable
  operators may remain undelegated.

## Testing

Tested on a CUDA GPU with TensorRT and ExecuTorch installed, building
Torch-TensorRT from source so the C++ runtime was active.

Verified by compiling a real model with TensorRT and then exporting it:

- static-shape `export()` produces a loadable `.pte`
- dynamic-shape `export()` produces a loadable `.pte`, matching the output of the
  previous one-step save path
- the source `GraphModule` still holds its original engine nodes afterward
- engine payloads are shared, not re-serialized
- `save(..., output_format="executorch")` still writes a `.pte`
- `save()` still raises `TypeError` for a non-list `partitioners` or
  `compile_specs`

Also ran the ExecuTorch test directory (76 passed). Two GPU tests in
`test_cuda_partitioner_composition.py` fail, and they fail the same way with this
change reverted, so they are not caused by it.

Static checks: formatting, import ordering, lint, and type checking on the changed
files.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: tests Issues re: Tests documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants